from core.text2image import text2imageN
from core.image2image import image2image, image2imageN
from core.image2video import image2video, clear
from core.convert import concatenate
from diffusers.utils import load_image
import streamlit as st
from typing import List
from datetime import datetime
import os

image_model_id = "../../hf/stable-diffusion-v1-5"
_processing_ = False

playv2_model_id = "playgroundai/playground-v2-1024px-aesthetic"
playv25_model_id = "playgroundai/playground-v2.5-1024px-aesthetic"
init_image = load_image("female2.png")  

def main():
  # gen_image_ui()
  # gen_video_with_images_ui()
  # test_image_to_image()
  # gen_image_with_images_ui()
  test_image_to_image()
  
  
def test_image_to_image():
  # a female cyborg
  # a female with  a pair of emerald earrings
  prompt = "a female with  a pair of emerald earrings" 
  negative_prompt = ""
  images = image2imageN(playv2_model_id, 1, prompt, negative_prompt,init_image)
  ts = datetime.now().strftime("%Y%m%d%H%M%S")
  output_dir = f"topic/{ts}"
  os.mkdir(output_dir) 
  for i, v in enumerate(images):
    v.save(f"{output_dir}/{i}.png")
  


def gen_image_with_images_ui():
  st.set_page_config(page_title="图->图", page_icon="图->图")
  st.title("图=>图")
  with st.form(key="my_form"):
    prompt = st.text_area(
      label="输入图片提示词..",
      help= "这里输入生成图片的提示词.",
      key="prompt",
      max_chars = 500,
    )
    negative_prompt = st.text_area(
      label="输入图片提示词..",
      help= "这里输入生成图片的negative提示词.",
      key="negative_prompt",
      max_chars = 1500,
    )
    submit_btn = st.form_submit_button(label="生成..")
    if prompt and submit_btn:
      global _processing_
      if _processing_:
        st.title('_服务器_ :red[正在处理其他请求] 请稍后尝试 :warning:')
        return
      _processing_ = True
      with st.spinner(text="正在生成图片..."):
        images = image2imageN(playv2_model_id, 3, prompt, negative_prompt, init_image)
        images.insert(0, init_image)
        st.subheader("您生成的图片")
        grid_image_ui(2, images)
        
def gen_video_with_images_ui():
  st.set_page_config(page_title="文->图->视频测试", page_icon="📖=>🎬")
  st.title("📖=>🎬")
  with st.form(key="my_form"):
    query_image = st.text_area(
      label="输入图片提示词..",
      help= "这里输入生成图片的提示词.",
      key="query_image",
      max_chars = 500,
    )
    query_video = st.text_area(
      label="输入视频提示词..",
      help= "这里输入生成视频的提示词.",
      key="query_video",
      max_chars = 500,
    )
    submit_btn = st.form_submit_button(label="生成..")
    
    if query_image and submit_btn:
      global _processing_
      if _processing_:
        st.title('_服务器_ :red[正在处理其他请求] 请稍后尝试 :warning:')
        return
      _processing_ = True
      if not query_video:
        query_video = query_image
      print(f"query_image is:<{query_image}>")
      print(f"query_video is:<{query_video}>")
      with st.spinner(text="正在生成视频..."):
        # prompt = "a photo of an astronaut riding a horse on mars"
        images = text2imageN(image_model_id, prompt = query_image, n=4)
        # images = [text2image(image_model_id, prompt) for _ in range(4)]
        # image = Image.open("astronaut_rides_horse_1.png")
        # image.save("astronaut_rides_horse.png")
      st.subheader("您生成的图片")
      grid_image_ui(2, images)
      
      with st.spinner(text="正在生成视频..."):
        ts = datetime.now().strftime("%Y%m%d%H%M%S")
        out_dir = f"tmp/{ts}"
        os.mkdir(out_dir) 
        videos = [image2video(query_video, image, f"{out_dir}/{idx}.mp4") for idx, image in enumerate(images)]
        video_file = f"{out_dir}/final.mp4"
        print("videos:", videos)
        concatenate(videos, video_file) 
        print("==final done:", video_file)
        clear()
      # A Labrador retriever dressed as a computer hacker, wearing a black hoodie, sits in front of a computer. As it rapidly clicks on the keyboard, the light of the computer screen reflects on its face, creating an atmosphere of concentration
      
      st.subheader("您生成的视频")
      st.video(data=video_file)
      _processing_ = False

def grid_image_ui(n_cols: int, images:List):
  n_rows = 1 + len(images) // int(n_cols)
  rows = [st.container() for _ in range(n_rows)]
  cols_per_row = [r.columns(n_cols) for r in rows]
  cols = [column for row in cols_per_row for column in row]

  for image_index, image in enumerate(images):
      cols[image_index].image(image)

if __name__ == "__main__":
  main()
 